from torch import nn
import torch

__all__ = ['AverageFusion']


class AverageFusion(nn.Module):

    def __init__(self, backbone, pretrained=True, hidden_size=1000):
        super(AverageFusion, self).__init__()
        self.rgb_stream = backbone(pretrained=pretrained, num_classes=hidden_size)
        self.msr_stream = backbone(pretrained=pretrained, num_classes=hidden_size)
        # self.q_kernel = nn.Parameter(torch.FloatTensor(1, hidden_size))
        # nn.init.normal_(self.q_kernel.weight, 0, 0.01)

    def forward(self, x_rgb, x_msr):
        x_rgb = self.rgb_stream(x_rgb)
        x_msr = self.msr_stream(x_msr)
        x_rgb_softmax = torch.softmax(x_rgb, dim=1)
        x_msr_softmax = torch.softmax(x_msr, dim=1)
        x = (x_msr_softmax + x_rgb_softmax) / 2.
        return x
